# generate_data_optimized.py
import numpy as np
import pandas as pd
import os
import matplotlib.pyplot as plt
from scipy import signal
from enum import Enum


class AnomalyType(Enum):
    """异常类型枚举"""
    NONE = "none"  # 无异常
    POINT = "point"  # 点异常
    SHIFT = "shift"  # 突变异常
    TREND_ACCEL = "trend_accel"  # 加速趋势异常
    TREND_REVERSE = "trend_reverse"  # 趋势反转异常
    PERIOD_MISSING = "period_missing"  # 周期缺失异常
    PERIOD_PHASE = "period_phase"  # 相位偏移异常
    CORRELATION = "correlation"  # 相关性破坏异常


def generate_dataset(output_dir, T=2000, anomaly_type=AnomalyType.NONE,
                     anomaly_params=None, seed=42):
    """
    生成多元时间序列数据集，支持参数化异常注入

    Parameters:
    - output_dir: 输出目录
    - T: 时间序列长度
    - anomaly_type: 异常类型
    - anomaly_params: 异常参数字典
    - seed: 随机种子
    """
    os.makedirs(output_dir, exist_ok=True)
    freq = "1min"
    start_time = "2025-10-31 00:00:00"
    np.random.seed(seed)

    timestamps = pd.date_range(start=start_time, periods=T, freq=freq)
    t = np.arange(T) / 100.0  # 归一化时间

    # 1. 基础信号生成
    print("Generating base signals...")

    # 传感器A: 主要趋势 + 季节性
    sensor_A = (
            2.0 * np.sin(0.1 * t) +  # 主要周期
            0.5 * np.sin(0.5 * t) +  # 高频周期
            0.1 * t +  # 趋势
            0.2 * np.random.randn(T)  # 噪声
    )

    # 传感器B: 与A相关但不同相位和振幅
    sensor_B = (
            1.5 * np.cos(0.1 * t + 0.5) +  # 相关但不同相位
            0.3 * np.sin(0.8 * t) +  # 额外周期
            0.05 * t +  # 较弱趋势
            0.15 * np.random.randn(T)  # 噪声
    )

    # 传感器C: 更复杂的模式
    sensor_C = (
            np.sin(0.08 * t) * np.cos(0.15 * t) +  # 调制信号
            0.3 * np.sin(1.2 * t) +  # 高频
            0.25 * np.random.randn(T)  # 噪声
    )

    # 2. 根据参数注入异常
    if anomaly_type != AnomalyType.NONE:
        print(f"Injecting {anomaly_type.value} anomaly...")

        # 使用默认参数或传入参数
        params = anomaly_params or get_default_anomaly_params(anomaly_type, T)

        # 注入特定类型的异常
        sensor_A, sensor_B, sensor_C = inject_anomaly(
            sensor_A, sensor_B, sensor_C, anomaly_type, params
        )

    # 3. 添加季节性上下文（基于时间的模式）
    print("Adding seasonal context...")
    hour_of_day = np.array([ts.hour for ts in timestamps])
    day_of_week = np.array([ts.dayofweek for ts in timestamps])

    # 日周期模式
    daily_pattern = 0.8 * np.sin(2 * np.pi * hour_of_day / 24)
    # 周周期模式
    weekly_pattern = 0.3 * np.sin(2 * np.pi * day_of_week / 7)

    sensor_A += daily_pattern + weekly_pattern
    sensor_B += daily_pattern * 0.6 + weekly_pattern * 0.4
    sensor_C += daily_pattern * 0.4 + weekly_pattern * 0.2

    # 4. 保存数据
    print("Saving data files...")
    sensors_data = {
        "sensor_A": sensor_A,
        "sensor_B": sensor_B,
        "sensor_C": sensor_C
    }

    for name, data in sensors_data.items():
        df = pd.DataFrame({
            "timestamp": timestamps,
            "value": data
        })
        df.to_csv(os.path.join(output_dir, f"{name}.csv"), index=False)
        print(f"Saved {name}.csv with {len(data)} data points")

    # 5. 创建标注文件
    if anomaly_type != AnomalyType.NONE:
        create_annotation_file(timestamps, anomaly_type, output_dir)

    # 6. 绘制图表
    plot_comprehensive_analysis(timestamps, sensors_data, output_dir,
                                anomaly_type != AnomalyType.NONE)

    print(f"\n✅ Dataset generation completed! Anomaly type: {anomaly_type.value}")
    print(f"📁 Output directory: {output_dir}")
    return sensors_data, timestamps


def get_default_anomaly_params(anomaly_type, T):
    """获取默认的异常参数"""
    base_params = {
        'start_position': T // 3,  # 异常开始位置
        'duration': min(200, T // 10),  # 异常持续时间
        'intensity': 1.0  # 异常强度
    }

    if anomaly_type == AnomalyType.POINT:
        return {
            **base_params,
            'positions': [T // 4, T // 2, T * 3 // 4],
            'magnitudes': [3.0, -2.0, 4.0],
            'sensors_affected': [0, 1, 2]  # 0:A, 1:B, 2:C
        }

    elif anomaly_type == AnomalyType.SHIFT:
        return {
            **base_params,
            'shift_magnitudes': [2.0, -1.5, 0.5],
            'sensors_affected': [0, 1, 2]
        }

    elif anomaly_type == AnomalyType.TREND_ACCEL:
        return {
            **base_params,
            'acceleration_factor': 0.01,
            'sensors_affected': [0, 1]
        }

    elif anomaly_type == AnomalyType.TREND_REVERSE:
        return {
            **base_params,
            'reversal_strength': -0.02,
            'sensors_affected': [0, 1, 2]
        }

    elif anomaly_type == AnomalyType.PERIOD_MISSING:
        return {
            **base_params,
            'missing_interval': 20,
            'suppression_strength': 1.5,
            'sensors_affected': [0, 1]
        }

    elif anomaly_type == AnomalyType.PERIOD_PHASE:
        return {
            **base_params,
            'phase_shift': np.pi / 2,
            'sensors_affected': [0, 1]
        }

    elif anomaly_type == AnomalyType.CORRELATION:
        return {
            **base_params,
            'corruption_strength': 2.0,
            'corruption_frequency': 0.5,
            'sensors_affected': [1]  # 主要影响传感器B
        }

    return base_params


def inject_anomaly(sensor_A, sensor_B, sensor_C, anomaly_type, params):
    """根据异常类型注入异常"""
    sensors = [sensor_A.copy(), sensor_B.copy(), sensor_C.copy()]
    T = len(sensor_A)

    start = params['start_position']
    duration = params['duration']
    end = min(start + duration, T)

    if anomaly_type == AnomalyType.POINT:
        inject_point_anomalies(sensors, params)

    elif anomaly_type == AnomalyType.SHIFT:
        inject_shift_anomaly(sensors, start, end, params)

    elif anomaly_type == AnomalyType.TREND_ACCEL:
        inject_trend_acceleration(sensors, start, end, params)

    elif anomaly_type == AnomalyType.TREND_REVERSE:
        inject_trend_reversal(sensors, start, end, params)

    elif anomaly_type == AnomalyType.PERIOD_MISSING:
        inject_period_missing(sensors, start, end, params)

    elif anomaly_type == AnomalyType.PERIOD_PHASE:
        inject_phase_shift(sensors, start, end, params)

    elif anomaly_type == AnomalyType.CORRELATION:
        inject_correlation_breakdown(sensors, start, end, params)

    return sensors[0], sensors[1], sensors[2]


def inject_point_anomalies(sensors, params):
    """注入点异常"""
    positions = params['positions']
    magnitudes = params['magnitudes']
    sensors_affected = params['sensors_affected']

    for pos, mag, sensor_idx in zip(positions, magnitudes, sensors_affected):
        if pos < len(sensors[sensor_idx]):
            sensors[sensor_idx][pos] += mag
            print(f"Point anomaly: Sensor {sensor_idx} at index {pos}, magnitude {mag}")


def inject_shift_anomaly(sensors, start, end, params):
    """注入突变异常"""
    magnitudes = params['shift_magnitudes']
    sensors_affected = params['sensors_affected']

    for sensor_idx, mag in zip(sensors_affected, magnitudes):
        sensors[sensor_idx][start:end] += mag
        print(f"Shift anomaly: Sensor {sensor_idx} from {start} to {end}, shift {mag}")


def inject_trend_acceleration(sensors, start, end, params):
    """注入加速趋势异常"""
    factor = params['acceleration_factor']
    sensors_affected = params['sensors_affected']
    time_window = np.arange(end - start)

    trend = factor * time_window ** 1.5
    for sensor_idx in sensors_affected:
        sensors[sensor_idx][start:end] += trend
        print(f"Trend acceleration: Sensor {sensor_idx} from {start} to {end}")


def inject_trend_reversal(sensors, start, end, params):
    """注入趋势反转异常"""
    strength = params['reversal_strength']
    sensors_affected = params['sensors_affected']
    time_window = np.arange(end - start)

    trend = strength * time_window
    for sensor_idx in sensors_affected:
        sensors[sensor_idx][start:end] += trend
        print(f"Trend reversal: Sensor {sensor_idx} from {start} to {end}")


def inject_period_missing(sensors, start, end, params):
    """注入周期缺失异常"""
    interval = params['missing_interval']
    strength = params['suppression_strength']
    sensors_affected = params['sensors_affected']

    for sensor_idx in sensors_affected:
        for i in range(start, end, interval):
            if i + 5 < end:
                sensors[sensor_idx][i:i + 5] -= strength
        print(f"Period missing: Sensor {sensor_idx} from {start} to {end}")


def inject_phase_shift(sensors, start, end, params):
    """注入相位偏移异常"""
    phase_shift = params['phase_shift']
    sensors_affected = params['sensors_affected']
    time_window = np.arange(end - start)

    shift_signal = 1.0 * np.sin(0.1 * (time_window + phase_shift / 0.1))
    for sensor_idx in sensors_affected:
        sensors[sensor_idx][start:end] += shift_signal
        print(f"Phase shift: Sensor {sensor_idx} from {start} to {end}")


def inject_correlation_breakdown(sensors, start, end, params):
    """注入相关性破坏异常"""
    strength = params['corruption_strength']
    frequency = params['corruption_frequency']
    sensors_affected = params['sensors_affected']
    time_window = np.arange(end - start)

    corruption = strength * np.sin(frequency * time_window)
    for sensor_idx in sensors_affected:
        sensors[sensor_idx][start:end] += corruption
        print(f"Correlation breakdown: Sensor {sensor_idx} from {start} to {end}")


def create_annotation_file(timestamps, anomaly_type, output_dir):
    """创建异常标注文件"""
    T = len(timestamps)
    params = get_default_anomaly_params(anomaly_type, T)
    start = params['start_position']
    duration = params['duration']
    end = min(start + duration, T)

    # 创建标注数组
    anomaly_label = np.zeros(T, dtype=int)
    anomaly_label[start:end] = 1

    annotation_df = pd.DataFrame({
        'timestamp': timestamps,
        'anomaly': anomaly_label,
        'anomaly_type': anomaly_type.value
    })

    annotation_df.to_csv(os.path.join(output_dir, "anomaly_annotations.csv"), index=False)

    # 创建异常摘要
    summary_df = pd.DataFrame({
        'anomaly_type': [anomaly_type.value],
        'start_index': [start],
        'end_index': [end - 1],
        'start_time': [timestamps[start]],
        'end_time': [timestamps[end - 1]],
        'duration': [duration]
    })
    summary_df.to_csv(os.path.join(output_dir, "anomaly_summary.csv"), index=False)

    print(f"✅ Annotation files created for {anomaly_type.value} anomaly")


def plot_comprehensive_analysis(timestamps, sensors_data, output_dir, has_anomalies):
    """绘制综合分析图表"""
    fig, axes = plt.subplots(3, 1, figsize=(16, 12))

    # 子图1: 原始数据
    axes[0].plot(timestamps, sensors_data['sensor_A'], label='Sensor A', linewidth=1, alpha=0.8)
    axes[0].plot(timestamps, sensors_data['sensor_B'], label='Sensor B', linewidth=1, alpha=0.8)
    axes[0].plot(timestamps, sensors_data['sensor_C'], label='Sensor C', linewidth=1, alpha=0.8)
    axes[0].set_title('Raw Sensor Data')
    axes[0].set_ylabel('Value')
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)

    # 子图2: 相关性分析
    window_size = min(50, len(timestamps) // 10)
    correlation_AB = pd.Series(sensors_data['sensor_A']).rolling(window_size).corr(pd.Series(sensors_data['sensor_B']))

    axes[1].plot(timestamps, correlation_AB, label='A-B Correlation', linewidth=1, color='green')
    axes[1].set_title(f'Rolling Correlation (window={window_size})')
    axes[1].set_ylabel('Correlation')
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)

    # 子图3: 异常标注
    if has_anomalies:
        try:
            annotation_path = os.path.join(output_dir, "anomaly_annotations.csv")
            if os.path.exists(annotation_path):
                annotation_df = pd.read_csv(annotation_path)
                axes[2].plot(timestamps, annotation_df['anomaly'],
                             color='red', linewidth=2, label='Anomaly Region', drawstyle='steps-post')
        except:
            pass

    axes[2].set_title('Anomaly Annotations')
    axes[2].set_xlabel('Timestamp')
    axes[2].set_ylabel('Anomaly (0/1)')
    axes[2].set_yticks([0, 1])
    axes[2].set_yticklabels(['Normal', 'Anomaly'])
    axes[2].legend()
    axes[2].grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "comprehensive_analysis.png"), dpi=150, bbox_inches='tight')
    plt.show()


def generate_all_anomaly_datasets():
    """生成所有异常类型的数据集"""
    base_dir = "anomaly_datasets"
    T = 2000

    # 生成无异常的训练数据
    print("=" * 60)
    print("Generating Clean Training Data")
    print("=" * 60)
    generate_dataset(
        os.path.join(base_dir, "train_clean"),
        T=T,
        anomaly_type=AnomalyType.NONE
    )

    # 生成各种异常类型的测试数据
    anomaly_types = [
        AnomalyType.POINT,
        AnomalyType.SHIFT,
        AnomalyType.TREND_ACCEL,
        AnomalyType.TREND_REVERSE,
        AnomalyType.PERIOD_MISSING,
        AnomalyType.PERIOD_PHASE,
        AnomalyType.CORRELATION
    ]

    for anomaly_type in anomaly_types:
        print("\n" + "=" * 60)
        print(f"Generating {anomaly_type.value} Anomaly Dataset")
        print("=" * 60)

        generate_dataset(
            os.path.join(base_dir, f"test_{anomaly_type.value}"),
            T=T,
            anomaly_type=anomaly_type
        )


def generate_specific_anomaly(anomaly_type, custom_params=None):
    """生成特定类型的异常数据"""
    output_dir = f"test_{anomaly_type.value}"
    T = 2000

    print(f"Generating {anomaly_type.value} anomaly dataset...")
    generate_dataset(
        output_dir,
        T=T,
        anomaly_type=anomaly_type,
        anomaly_params=custom_params
    )


if __name__ == "__main__":
    # 方法1: 生成所有异常类型的数据集
    # generate_all_anomaly_datasets()

    # 方法2: 生成特定类型的异常数据（取消注释使用）
    generate_specific_anomaly(AnomalyType.PERIOD_PHASE)

    # 方法3: 使用自定义参数生成异常数据（取消注释使用）
    # custom_params = {
    #     'start_position': 500,
    #     'duration': 300,
    #     'magnitudes': [5.0, -3.0, 2.0],
    #     'sensors_affected': [0, 1, 2]
    # }
    # generate_specific_anomaly(AnomalyType.POINT, custom_params)

    print("\n🎉 All anomaly datasets generated successfully!")
    print("📁 Dataset structure:")
    print("   - anomaly_datasets/train_clean/ (clean training data)")
    print("   - anomaly_datasets/test_point/ (point anomalies)")
    print("   - anomaly_datasets/test_shift/ (shift anomalies)")
    print("   - anomaly_datasets/test_trend_accel/ (trend acceleration)")
    print("   - anomaly_datasets/test_trend_reverse/ (trend reversal)")
    print("   - anomaly_datasets/test_period_missing/ (period missing)")
    print("   - anomaly_datasets/test_period_phase/ (phase shift)")
    print("   - anomaly_datasets/test_correlation/ (correlation breakdown)")